package com.mings.ai.service;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xssf.extractor.XSSFExcelExtractor;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.hssf.extractor.ExcelExtractor;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.InputStream;

@Service
public class DocumentProcessor {

    /**
     * 处理多种格式的文档，提取文本内容
     */
    public String extractTextFromFile(MultipartFile file) throws Exception {
        String fileName = file.getOriginalFilename();
        String extension = getFileExtension(fileName);
        
        try (InputStream inputStream = file.getInputStream()) {
            switch (extension.toLowerCase()) {
                case "pdf":
                    return extractTextFromPdf(inputStream);
                case "docx":
                    return extractTextFromDocx(inputStream);
                case "doc":
                    return extractTextFromDoc(inputStream);
                case "xlsx":
                    return extractTextFromXlsx(inputStream);
                case "xls":
                    return extractTextFromXls(inputStream);
                case "txt":
                    return new String(file.getBytes(), "UTF-8");
                default:
                    throw new UnsupportedOperationException("不支持的文件格式: " + extension);
            }
        }
    }

    /**
     * 提取PDF文件文本
     */
    private String extractTextFromPdf(InputStream inputStream) throws Exception {
        try (PDDocument document = PDDocument.load(inputStream)) {
            PDFTextStripper stripper = new PDFTextStripper();
            return stripper.getText(document);
        }
    }

    /**
     * 提取DOCX文件文本
     */
    private String extractTextFromDocx(InputStream inputStream) throws Exception {
        try (XWPFDocument document = new XWPFDocument(inputStream)) {
            XWPFWordExtractor extractor = new XWPFWordExtractor(document);
            return extractor.getText();
        }
    }

    /**
     * 提取DOC文件文本
     */
    private String extractTextFromDoc(InputStream inputStream) throws Exception {
        try (HWPFDocument document = new HWPFDocument(inputStream)) {
            WordExtractor extractor = new WordExtractor(document);
            return extractor.getText();
        }
    }

    /**
     * 提取XLSX文件文本
     */
    private String extractTextFromXlsx(InputStream inputStream) throws Exception {
        try (XSSFWorkbook workbook = new XSSFWorkbook(inputStream)) {
            XSSFExcelExtractor extractor = new XSSFExcelExtractor(workbook);
            extractor.setFormulasNotResults(true);
            extractor.setIncludeSheetNames(true);
            return extractor.getText();
        }
    }

    /**
     * 提取XLS文件文本
     */
    private String extractTextFromXls(InputStream inputStream) throws Exception {
        try (HSSFWorkbook workbook = new HSSFWorkbook(inputStream)) {
            ExcelExtractor extractor = new ExcelExtractor(workbook);
            extractor.setFormulasNotResults(true);
            extractor.setIncludeSheetNames(true);
            return extractor.getText();
        }
    }

    /**
     * 获取文件扩展名
     */
    private String getFileExtension(String fileName) {
        if (fileName == null || !fileName.contains(".")) {
            return "";
        }
        return fileName.substring(fileName.lastIndexOf(".") + 1);
    }

    /**
     * 检查是否支持的文件格式
     */
    public boolean isSupportedFormat(String fileName) {
        String extension = getFileExtension(fileName);
        return extension.matches("(?i)(pdf|docx|doc|xlsx|xls|txt)");
    }
}